import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 步骤一（替换sans-serif字体）
plt.rcParams['axes.unicode_minus'] = False  # 步骤二（解决坐标轴负数的负号显示问题）

# 高斯概率模型的分布估计算法类
class EDA:
    def __init__(self):
        self.learning_rate = 0.01
        self.iteration_num = 200
        self.N = 1000  # 种群大小
        self.M = None  # 所有个体集合
        self.encode_x = None  # 编码值x
        self.encode_y = None  # 编码值y
        self.decode_x = np.zeros((self.N, 1))  # 解码值x
        self.decode_y = np.zeros((self.N, 1))  # 解码值y
        self.fitness = None  # 适应度
        self.dominant_population = None  # 优势种群
        self.miu = None
        self.sigma = None
        self.best = [0]
        self.best_x = None
        self.best_y = None
        self.init_population()  # 初始化产生种群
        self.init_guass()  # 初始化概率模型
        self.x = None
        self.y = None

    # 随机产生种群
    def init_population(self):
        self.M = np.random.randint(2, size=(self.N, 20))

    # 生成坐标点
    def generate_xy(self):
        self.encode_x = self.M[:, 0:10]
        self.encode_y = self.M[:, 10:20]

    # 目标函数
    def Camel(self, x, y):
        camel = (4 - 2.1 * x ** 2 + x ** 4 / 3) * x ** 2 + x * y + (-4 + 4 * y ** 2) * y ** 2
        return camel

    # 解码
    def decode(self):
        for i in range(self.N):
            xt = 0
            yt = 0
            for j in range(10):
                xt += self.encode_x[i, j] * 2 ** (-(j - 9))
                yt += self.encode_y[i, j] * 2 ** (-(j - 9))
            self.decode_x[i] = 6 * xt / (1024 - 1) - 3
            self.decode_y[i] = 4 * yt / (1024 - 1) - 2

    # 计算单个点的函数值
    def decode_xy(self, best):
        self.x = best[:, 0:10]
        self.y = best[:, 10:20]
        xt = 0
        yt = 0
        for j in range(10):
            xt += self.x[0, j] * 2 ** (-(j - 9))
            yt += self.y[0, j] * 2 ** (-(j - 9))
        x = 6 * xt / (1024 - 1) - 3
        y = 4 * yt / (1024 - 1) - 2
        camel = (4 - 2.1 * x ** 2 + x ** 4 / 3) * x ** 2 + x * y + (-4 + 4 * y ** 2) * y ** 2
        return x, y, camel

    # 适应度函数
    def fitness_function(self, m):
        Cmax = np.max(m)
        self.fitness = Cmax - m

    # 获得优势种群
    def get_dominant(self):
        # 按照从大到小排列,返回索引值
        num = np.argsort(-self.fitness)
        # 取前500个优势个体索引值
        num = num[100:]
        m = []
        # 获得优势种群
        for i in num:
            m.append(self.M[i])
        self.dominant_population = np.asarray(m)

    # 获取最优、次优、最差个体
    def get_best1_best2_worst(self):
        best1 = self.dominant_population[0, :]
        best2 = self.dominant_population[1, :]
        wrost = self.dominant_population[-1, :]
        return best1, best2, wrost

    # 初始化高斯概率模型
    def init_guass(self):
        # 每一列的平均值,生成一个行向量
        self.miu = np.mean(self.M, axis=0)
        self.sigma = np.sqrt(np.mean((self.M - self.miu) ** 2, axis=0))

    # 更新高斯概率模型
    def update_guass(self, best1, best2, wrost):
        # 更新u
        self.miu = (1 - self.learning_rate) * self.miu + self.learning_rate * (best1 + best2 - wrost)
        # 更新a
        average = np.mean(self.dominant_population, axis=0)  # 计算每一列的均值
        self.sigma = (1 - self.learning_rate) * self.sigma + self.learning_rate * np.sqrt(
            np.mean((self.dominant_population - average) ** 2, axis=0))

    # 产生新种群
    def generate_population(self):
        # 产生高斯分布的样本
        self.M = np.random.normal(loc=self.miu, scale=self.sigma, size=(self.N, 20))
        # 设置阈值
        for i in range(len(self.M)):
            for j in range(20):
                if self.M[i, j] >= 0.5: self.M[i, j] = 1
                if self.M[i, j] <  0.5: self.M[i, j] = 0

    # 获取最优值
    def get_best(self, fx):
        best = np.min(fx)
        fx_min = np.argmin(fx)
        x = self.decode_x[fx_min]
        y = self.decode_y[fx_min]

        if best < np.min(self.best):
            self.best.append(best)
            self.best_x = x
            self.best_y = y
        else:
            self.best.append(np.min(self.best))


    # 解析
    def calculate(self):
        # 计算x、y
        self.generate_xy()
        # 解码
        self.decode()
        # 计算函数值
        fx = self.Camel(self.decode_x, self.decode_y)
        # 获得最优值
        self.get_best(fx)

        # 获取适应度
        self.fitness_function(fx)
        # 获取优势种群
        self.get_dominant()
        # 获取最优、次优、最差个体
        best1, best2, wrost = self.get_best1_best2_worst()
        # 更新高斯概率模型
        self.update_guass(best1, best2, wrost)
        # 产生样本
        self.generate_population()

    # 迭代函数
    def iteration(self):
        for i in range(self.iteration_num):
            self.calculate()
            print('x:', self.best_x,'y:',self.best_y,'最优值:',self.best[-1])


a = EDA()
a.iteration()
fig = plt.figure()
plt.xlabel('迭代次数')
plt.ylabel('最优函数值')
plt.plot(np.linspace(0, a.iteration_num, a.iteration_num), a.best[1:], color="black", label='函数值')
plt.legend()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
xm = np.linspace(-3, 3, 100)
ym = np.linspace(-2, 2, 100)
xm, ym = np.meshgrid(xm, ym)
zm = a.Camel(xm, xm)
# ax.scatter(a.decode_x, a.decode_y, np.min(a.best), color='r', s=20)
ax.plot_wireframe(xm, ym, zm, rstride=5, cstride=5, cmap=plt.cm.coolwarm)
plt.show()